function [ subjectinfo ] = buildDatasets( subjectinfo )
%BUILDDATASETS Imports, preprocesses and stores the raw ECoG data

% STEP 1
subjectinfo = prepareFilepath(subjectinfo); % Prepare the filepath structure
[subjectinfo datasets] = importRawDatasets(subjectinfo); % Import the raw data
subjectinfo = storeMarkersChannels(subjectinfo, datasets); % Store information on markers and channels

% STEP 2
[subjectinfo success1] = updateElectrodeLocations(subjectinfo);
[subjectinfo success2] = checkSetMarkerTransform(subjectinfo);
[subjectinfo success3] = checkSetRejectChannels(subjectinfo);
if ~(success1 && success2 && success3) % Terminate the process if not ready
    disp('Please make sure above tasks are completed, then run buildDatasets_Step2 again.');
    error('Script aborted.')
end

% STEP 3
subjectinfo = transformAndExportDatasets(subjectinfo); % Old buildDatasets_Step3.m
end

%% Preparation
function [ s ] = prepareFilepath( s )
if ~exist([s.FilepathBase 'Datasets\'], 'dir')
    mkdir(s.FilepathBase, 'Datasets');
end
end

%% Importing
function [ s datasets ] = importRawDatasets( s )
% Determine the number of destination datasets and start importing
tic; % Because of a bug in pop_load_trc...
datasets = unique([s.OriginalData.DestinationDataset]);
for d = datasets
    if ~isa(s.Datasets(d), 'DatasetInfo')
        error('Invalid or missing Dataset Info, update your subject information.');
    end
    setCount = 0;
    for i = 1:length(s.OriginalData)
        if s.OriginalData(i).DestinationDataset == d
            setCount = setCount + 1;
            EEG = pop_load_trc(s.OriginalData(i).Filename, '+v' );
            % In case this is the first datafile, use it as is
            if setCount == 1
                FINAL = EEG;
                % In case this is a following datafiles, merge them with the first
            elseif setCount > 1
                FINAL = pop_mergeset( FINAL, EEG, 0);
            end
        end
    end
    % If there is any data, save it to disk as unprocessed data
    if setCount > 0
        FINAL = pop_editset(FINAL, 'setname', [s.Name ' - Set ' num2str(d) ' Unprocessed']);
        FINAL = eeg_checkset( FINAL );
        % Store the information in subjectinfo
        s.Datasets(d).ImportedDataFilename = [s.FilepathBase 'Datasets\' s.PrimaryPrefix num2str(d) 'Unprocessed.set'];
        pop_saveset(FINAL, 'filename', s.Datasets(d).ImportedDataFilename);
    end
end
clear FINAL;
clear EEG;

% Check the srate (with Subject CW, srate is incorrectly set to 1)
srateOk = true;
for d = datasets
    EEG = pop_loadset( 'filename', s.Datasets(d).ImportedDataFilename);
    if (EEG.srate ~= 512)
        warning(['Dataset (' s.Datasets(d).ImportedDataFilename ') had an incorrect sampling rate after the import and will be reset to 512Hz']);
        EEG = pop_editset(EEG, 'srate', 512);
        EEG = eeg_checkset(EEG);
        pop_saveset(EEG, 'filename', s.Datasets(d).ImportedDataFilename);
    end
end
end

%% Find and store marker and channel information
function [ s ] = storeMarkersChannels( s, datasets )
x(length(datasets)).nbchan = 0; % preallocate size..
for d = datasets
    EEG = pop_loadset( 'filename', s.Datasets(d).ImportedDataFilename);
    % Create the marker transform
    s.Datasets(d).MarkerTransform = ...
        MarkerTransformInfo.CreateMarkerTransformInfo(EEG.event, EEG.srate);
    % Read the channel location info
    x(d).nbchan = EEG.nbchan;
    x(d).chanlocs = EEG.chanlocs;
    clear EEG;
end

% Compare the labels for all datasets making sure they are consistent.
for d = 2:length(x)
    if x(d).nbchan ~= x(1).nbchan
        error('Inconsistent number of channels between datasets.');
    end
    if ~min(strcmp({x(1).chanlocs(:).labels},{x(2).chanlocs(:).labels})) % some labels do not match..
        error('Inconsistent electrode labels between datasets.');
    end
end
% Get the label info and store it..
s.Electrodes = ElectrodeInfo.CreateElectrodeInfo(x(1).chanlocs);
end

%% Update electrode locations
function [ s readyToProceed ] = updateElectrodeLocations( s )
readyToProceed = true;
warning('Not implemented yet!');
disp(' --> Add channel locations in SubjectInfo.Electrodes, this isn''t automated yet..');
end

%% Set marker transforms
function [ s readyToProceed ] = checkSetMarkerTransform( s )
readyToProceed = true;
if exist([s.FilepathBase 'markersImport.mat'], 'file')
    % Import the transformation
    load([s.FilepathBase 'markersImport.mat']);
    for d = 1:length(s.Datasets)
        s.Datasets(d).MarkerTransform = ...
            MarkerTransformInfo.ImportMatrix( s.Datasets(d).MarkerTransform, ...
            markers.(['dataset' num2str(d)]) );
    end
    disp(['Markers loaded from ' s.FilepathBase 'markersImport.mat.']);
else
    % Export the current set of markers to file
    readyToProceed = false;
    for d = 1:length(s.Datasets)
        markers.(['dataset' num2str(d)]) = ...
            MarkerTransformInfo.ExportMarkerTransform(s.Datasets(d).MarkerTransform);
    end
    save ([s.FilepathBase 'markersExport.mat'], 'markers');
    disp(['Markers exported to ' s.FilepathBase 'markersExport.mat.']);
    disp(' --> Edit the variable in markersExport.mat and save the variable to markersImport.mat.');
end
end

%% Set Reject electrode channels
function [ s readyToProceed ] = checkSetRejectChannels( s )
readyToProceed = true;
if exist([s.FilepathBase 'rejectedChannels.mat'], 'file')
    load([s.FilepathBase 'rejectedChannels.mat']);
    for i = 1:length(rejectedChannels)
        s.Electrodes(strcmp({s.Electrodes(:).Label}, rejectedChannels(i))).RejectInPreprocessing = true;
    end
    disp('Channels marked as rejected: (will not reset/show previously rejected channels)');
    disp(rejectedChannels');
else
    readyToProceed = false;
    disp(' --> Save the electrode labels that should be rejected like:');
    disp('    rejectedChannels = {''R1+'' ''R2+'' ''ECG+'' ...}');
    disp('to a file named rejectedChannels.mat.');
end
end

%% Transform the data according to what is set in 'step 2' and save the files to disk
function [s] = transformAndExportDatasets(s)


% Preprocess each dataset as described in the subject information
for d = 1:length(s.Datasets)
    
    EEG = pop_loadset( 'filename', s.Datasets(d).ImportedDataFilename);
    
    % Apply the marker transform\
    disp('Applying the marker transformation...');
    disp(['  --> There are currently ' num2str(length(EEG.event)) ' marker(s) in dataset ' num2str(d) '.']);
    newcounter = 0;
    clear newevent;
    for i = 1:length(s.Datasets(d).MarkerTransform)
%         fprintf(1, 'Processing event %d of %d..\r', i, length(EEG.event))
        % Get the original EEG.event
        mt = s.Datasets(d).MarkerTransform(i);
        lat = mt.LatencyFrames;
        event = EEG.event([EEG.event(:).latency]==lat);
        %         event = EEG.event(i);
        
        switch mt.Action
            case 'None'
                % Just add to the new event list unchanged
                newcounter = newcounter + 1;
                newevent(newcounter) = event;
                disp(['Marker ' num2str(i) ' kept unchanged.']);
            case 'Rename'
                if ischar(mt.NewType) && ~isempty(mt.NewType)
                    % Change the event type, then add to list
                    event.type = mt.NewType;
                    newcounter = newcounter + 1;
                    newevent(newcounter) = event;
                    disp(['Renamed marker ' num2str(i) '.']);
                end
            case 'Delete'
                % Don't add to the list
                disp(['Deleted marker ' num2str(i) '.']);
            otherwise
                error('Unknown marker transform action.')
        end
    end
    EEG.event = newevent;
    disp(['  --> There are currently ' num2str(length(EEG.event)) ' marker(s) in dataset ' num2str(d) '.']);
    
    % Update the electrode locations
    disp('Updating the electrode locations...');
    for k=1:EEG.nbchan
        e = ElectrodeInfo.FindByLabel(s.Electrodes, EEG.chanlocs(k).labels);
        EEG.chanlocs(k).X = e.LocationXYZ(1);
        EEG.chanlocs(k).Y = e.LocationXYZ(2);
        EEG.chanlocs(k).Z = e.LocationXYZ(3);
    end
    
    % Reject indicated channels
    selectSet = find([s.Electrodes(:).RejectInPreprocessing] == false);
    selectChannels = {s.Electrodes(selectSet).Label};
    EEG = pop_select(EEG, 'channel', selectChannels);
    
    % Filter the data
    fig = figure;
    pop_spectopo(EEG, 1, [0 EEG.xmax], 'EEG' , 'percent', 15, 'freqrange', [0 250], 'electrodes', 'off');
    saveFigure(fig, [s.FilepathBase 'Datasets\'], [s.PrimaryPrefix num2str(d) ' (Before filtering)']);
    close(fig);
    EEG = pop_eegfilt( EEG, 2, 0, [], 0, 0); % Highpass FIR-filter
    EEG = pop_eegfilt( EEG, 0, 200, [], 0, 0); % Lowpass FIR-filter
    EEG = pop_eegfilt( EEG, 48.5, 51.5, [], 1, 1); % 3Hz Power line noise notch FFT-filter (sharper than FIR)
    EEG = pop_eegfilt( EEG, 148.5, 151.5, [], 1, 1); % 3Hz Power line noise notch FFT-filter (sharper than FIR)
    fig = figure;
    pop_spectopo(EEG, 1, [0 EEG.xmax], 'EEG' , 'percent', 15, 'freqrange', [0 250], 'electrodes', 'off');
    saveFigure(fig, [s.FilepathBase 'Datasets\'], [s.PrimaryPrefix num2str(d) ' (After filtering)']);
    close(fig);
    
    % Rereference the data
    EEG = pop_reref(EEG, []);
    
    % Save the dataset
    EEG = pop_editset(EEG, 'setname', [s.Name ' - Set ' num2str(d)]);
    EEG = eeg_checkset( EEG );
    
    s.Datasets(d).PreprocessedFilename = [s.FilepathBase 'Datasets\' s.PrimaryPrefix num2str(d) '.set'];
    EEG = pop_saveset(EEG, 'filename', s.Datasets(d).PreprocessedFilename);
    
end

end